package com.mentalhealthplatform.utils.decryptUtils.field;

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;
import java.util.Base64;

public class AESUtils {

    // 使用 AES CBC 模式，PKCS5 填充
    private static final String ALGORITHM = "AES/CBC/PKCS5Padding";

    /**
     * 确保密钥的长度为指定的字节数（例如：16 字节）
     * @param secretKey - 原始密钥
     * @param length - 目标长度（以字节为单位）
     * @return - 填充或截断后的密钥
     */
    private static byte[] normalizeKeyLength(String secretKey, int length) {
        byte[] keyBytes = secretKey.getBytes(StandardCharsets.UTF_8);

        // 如果密钥太短，使用 '0' 填充
        if (keyBytes.length < length) {
            byte[] newKey = new byte[length];
            System.arraycopy(keyBytes, 0, newKey, 0, keyBytes.length);
            return newKey;
        }

        // 如果密钥太长，截断到指定长度
        if (keyBytes.length > length) {
            byte[] truncatedKey = new byte[length];
            System.arraycopy(keyBytes, 0, truncatedKey, 0, length);
            return truncatedKey;
        }

        return keyBytes;
    }

    /**
     * 解密方法
     * @param encryptedText Base64 编码的加密数据（包含 IV 和密文）
     * @param secretKey 解密的密钥（16 字节）
     * @return 解密后的原始文本
     * @throws Exception 解密异常
     */
    public static String decrypt(String encryptedText, String secretKey) throws Exception {
        // 首先解码 Base64 加密数据，得到密文
        byte[] encryptedData = Base64.getDecoder().decode(encryptedText);

        // 提取前 16 字节作为 IV
        byte[] iv = new byte[16];
        System.arraycopy(encryptedData, 0, iv, 0, iv.length);
        IvParameterSpec ivParams = new IvParameterSpec(iv);

        // 获取密文数据
        byte[] cipherText = new byte[encryptedData.length - iv.length];
        System.arraycopy(encryptedData, iv.length, cipherText, 0, cipherText.length);

        // 确保密钥长度为 16 字节（AES-128）
        byte[] key = normalizeKeyLength(secretKey, 16);
        SecretKeySpec secretKeySpec = new SecretKeySpec(key, "AES");

        // 使用 CBC 模式和 IV 解密密文
        Cipher cipher = Cipher.getInstance(ALGORITHM);
        cipher.init(Cipher.DECRYPT_MODE, secretKeySpec, ivParams);

        // 解密数据
        byte[] decryptedData = cipher.doFinal(cipherText);

        // 返回解密后的字符串
        return new String(decryptedData, StandardCharsets.UTF_8);
    }

    /**
     * 加密方法（将 IV 附加到密文前）
     * @param plainText 明文文本
     * @param secretKey 加密的密钥（16 字节）
     * @return Base64 编码的加密数据（包含 IV 和密文）
     * @throws Exception 加密异常
     */
    public static String encrypt(String plainText, String secretKey) throws Exception {
        // 确保密钥长度为 16 字节（AES-128）
        byte[] key = normalizeKeyLength(secretKey, 16);
        SecretKeySpec secretKeySpec = new SecretKeySpec(key, "AES");

        // 随机生成 IV
        byte[] iv = new byte[16];
        SecureRandom random = new SecureRandom();
        random.nextBytes(iv);
        IvParameterSpec ivParams = new IvParameterSpec(iv);

        // 使用 CBC 模式加密
        Cipher cipher = Cipher.getInstance(ALGORITHM);
        cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec, ivParams);

        // 加密数据
        byte[] encryptedData = cipher.doFinal(plainText.getBytes(StandardCharsets.UTF_8));

        // 将 IV 和密文合并，并进行 Base64 编码
        byte[] combinedData = new byte[iv.length + encryptedData.length];
        System.arraycopy(iv, 0, combinedData, 0, iv.length);
        System.arraycopy(encryptedData, 0, combinedData, iv.length, encryptedData.length);

        return Base64.getEncoder().encodeToString(combinedData);
    }

    public static void main(String[] args) {
        try {
            String secretKey = "mysecretkey12345";  // 16 字节密钥
            String plainText = "Hello, AES Encryption!";  // 明文

            // 调用加密方法
            String encryptedText = encrypt(plainText, secretKey);
            System.out.println("Encrypted Text: " + encryptedText);

            // 调用解密方法
            String decryptedText = decrypt("ICxBAWZh4z+BLgnXMjRwWorwKilJiFWcfKn4ANCcrpv3AHoWE1lOOrkhHl52MzKY", secretKey);
            System.out.println("Decrypted Text: " + decryptedText);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
